#include "ioctl_channel.h"
#include "common/log.h"

#define IOCTL_MSG_ASYNC CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_READ_DATA | FILE_WRITE_DATA)
#define IOCTL_MSG_SYNC CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_BUFFERED, FILE_READ_DATA | FILE_WRITE_DATA)

namespace {
constexpr int DEVICE_ID_START = 0;
constexpr int DEVICE_ID_END = 16;

const std::wstring FILTER_RECV_DEVICE_NAME = L"\\\\.\\process_filter_recv_{}";
const std::wstring FILTER_WRITE_DEVICE_NAME = L"\\\\.\\process_filter_send_{}";

BOOL WINAPI ConsoleHandler(DWORD dwCtrlType)
{
    switch (dwCtrlType) {
        case CTRL_CLOSE_EVENT:
            IoctlChannel::GetInstance()->Disconnect();
            return TRUE;
        default:
            return FALSE;
    }
}
}  // namespace

IoctlChannel::~IoctlChannel()
{
    Disconnect();
}

HANDLE IoctlChannel::Connect(const std::wstring &deviceName)
{
    for (int i = DEVICE_ID_START; i <= DEVICE_ID_END; i++) {
        HANDLE device = CreateFileW(fmt::format(deviceName, i).c_str(),
            GENERIC_READ | GENERIC_WRITE,
            0,
            nullptr,
            OPEN_EXISTING,
            FILE_ATTRIBUTE_SYSTEM | FILE_FLAG_OVERLAPPED,
            nullptr);
        if (device == INVALID_HANDLE_VALUE) {
            LOG_ERR("failed to open device.");
            continue;
        }

        IO_MSG msg{};
        msg.FileName[0] = 1;
        msg.NameMaxLength = MSG_NAME_LENGTH;
        BOOL result = DeviceIoControl(device, IOCTL_MSG_SYNC, &msg, sizeof(msg), &msg, sizeof(msg), nullptr, nullptr);
        if (!result || msg.Message != MSG_STATE_DEVICE_OK) {
            LOG_ERR("check device failed.");
            continue;
        }
        return device;
    }
    return nullptr;
}

bool IoctlChannel::Connect()
{
    if (SetConsoleCtrlHandler(ConsoleHandler, TRUE) == 0) {
        LOG_ERR("error setting control handler");
        return false;
    }
    if ((devices[0] = Connect(FILTER_WRITE_DEVICE_NAME)) == nullptr) {
        LOG_ERR("connect recv device failed, err: {}", GetLastError());
        return false;
    }
    if ((devices[1] = Connect(FILTER_RECV_DEVICE_NAME)) == nullptr) {
        LOG_ERR("connect write device failed, err: {}", GetLastError());
        return false;
    }
    LOG_INFO("connect device success.");
    return true;
}

void IoctlChannel::Disconnect()
{
    LOG_INFO("disconnect.");
    CloseHandle(devices[0]);
    CloseHandle(devices[1]);
}

bool IoctlChannel::Reply(IO_MSG &msg)
{
    if (devices[0] == nullptr) {
        LOG_ERR("device is not ready.");
        return false;
    }
    std::lock_guard lock(sendLock);
    LOG_DEBUG("reply msg for driver.");
    return DeviceIoControl(devices[0], IOCTL_MSG_SYNC, &msg, sizeof(msg), &msg, sizeof(msg), nullptr, nullptr);
}

bool IoctlChannel::Recv(IO_MSG &msg)
{
    if (devices[1] == nullptr) {
        LOG_ERR("device is not ready.");
        return false;
    }
    LOG_DEBUG("recv msg from driver.");
    OVERLAPPED overLapped;
    overLapped.hEvent = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    if (!overLapped.hEvent) {
        LOG_ERR("Failed to create event for overlap while recv");
        return false;
    }
    auto ret =
        DeviceIoControl(devices[1], IOCTL_MSG_ASYNC, &msg, sizeof(IO_MSG), &msg, sizeof(IO_MSG), nullptr, &overLapped);
    if (!ret && GetLastError() != ERROR_IO_PENDING) {
        LOG_ERR("DeviceIoControl failed. Error: {}", GetLastError());
        CloseHandle(overLapped.hEvent);
        return false;
    }
    ret = WaitForSingleObject(overLapped.hEvent, INFINITE) == WAIT_OBJECT_0;
    CloseHandle(overLapped.hEvent);
    return ret;
}

std::shared_ptr<IoctlChannel> IoctlChannel::GetInstance()
{
    static std::shared_ptr<IoctlChannel> channel = std::make_shared<IoctlChannel>();
    return channel;
}